/*******************************************************************************
 * Copyright (c) 2008-2009 zhang yuexiang.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *      http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */ 
package org.xfeep.alaya.jta;

import static org.xfeep.asura.core.reflect.TypeInfoUtil.handleThrowable;

import java.io.ByteArrayInputStream;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.lang.reflect.Modifier;
import java.security.ProtectionDomain;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;

import javax.ejb.TransactionAttribute;
import javax.ejb.TransactionAttributeType;
import javax.ejb.TransactionManagement;
import javax.ejb.TransactionManagementType;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtField;
import javassist.CtMethod;
import javassist.CtNewMethod;
import javassist.CtPrimitiveType;
import javassist.LoaderClassPath;
import javassist.NotFoundException;
import javassist.bytecode.AnnotationsAttribute;
import javassist.bytecode.ClassFile;
import javassist.bytecode.ConstPool;
import javassist.bytecode.FieldInfo;
import javassist.bytecode.annotation.Annotation;

import org.xfeep.alaya.jta.agent.PreMainAgent;
import org.xfeep.asura.core.annotation.Activate;
import org.xfeep.asura.core.annotation.Deactivate;
import org.xfeep.asura.core.annotation.Ref;
import org.xfeep.asura.core.annotation.Service;
import org.xfeep.asura.core.logger.LoggerService;

@Service
public class WeaverService implements ClassFileTransformer {
	
	final static ClassLoader selfLoader = WeaverService.class.getClassLoader();
	
	@Ref
	LoggerService loggerService;
	
	ConcurrentHashMap<ClassLoader, Future<ClassPool>> classPoolMap = new ConcurrentHashMap<ClassLoader, Future<ClassPool>>(); 

	public ClassPool getClassPool(final ClassLoader loader) {
		
		//use a concurrent cache implementation suggested by Java Concurrency in Practice
		//5.6. Building an Efficient, Scalable Result Cache
		while(true) {
			Future<ClassPool> f = classPoolMap.get(loader);
	        if (f == null) {
	            Callable<ClassPool> builder = new Callable<ClassPool>() {
	                public ClassPool call()  {
	                	ClassPool rt = new ClassPool();
	                	rt.appendClassPath( new LoaderClassPath(loader) );
	                	return rt;
	                }
	            };
	            FutureTask<ClassPool> ft = new FutureTask<ClassPool>(builder);
	            f = classPoolMap.putIfAbsent(loader, ft);
	            if (f == null) { 
	            	f = ft; ft.run(); 
	            }
	        }
	        try {
	            return f.get();
	        } catch (CancellationException e) {
	        	classPoolMap.remove(loader, f);
	        } catch (ExecutionException e) {
	            throw handleThrowable(e.getCause());
	        } catch (InterruptedException e) {
			}
		}
		
	}
	
	public static <T> T getAnnotation(CtClass cls, Class<T> an) throws ClassNotFoundException{
		for (Object o : cls.getAvailableAnnotations()){
			if (an.isInstance(o)){
					return an.cast(o);
			}
		}
		return null;
	}
	
	public static <T> T getAnnotation(CtMethod m, Class<T> an) throws ClassNotFoundException{
		for (Object o : m.getAvailableAnnotations()){
			if (an.isInstance(o)){
					return an.cast(o);
			}
		}
		return null;
	}
	
	@SuppressWarnings("unchecked")
	public static CtClass addRef(ClassPool classPool, CtClass c, String fieldType, String field) throws CannotCompileException, NotFoundException{
		
		c.addField(new CtField(classPool.get(fieldType), field, c ));
		ClassFile frontClsFile = c.getClassFile();
		frontClsFile.setVersionToJava5();
		ConstPool cp = frontClsFile.getConstPool();
		FieldInfo fi = null;
		for (FieldInfo fie : (List<FieldInfo>)frontClsFile.getFields()){
			if (fie.getName().equals(field)){
				fi = fie;
				break;
			}
		}
		AnnotationsAttribute fattr = new AnnotationsAttribute(cp, AnnotationsAttribute.visibleTag);
		Annotation fesRef = new Annotation(Ref.class.getName(), cp);
		fattr.setAnnotation(fesRef);
		fi.addAttribute(fattr);
		return c;
	}
	
	
	final public byte[] transform(ClassLoader loader, String className,
			Class<?> classBeingRedefined, ProtectionDomain protectionDomain,
			byte[] classfileBuffer) throws IllegalClassFormatException {
		
		while (loader != null){
			if (loader == selfLoader){
				break;
			}
			loader = loader.getParent();
		}
		if (loader == null ){
			return classfileBuffer;
		}
		
		byte[] rt = classfileBuffer;
		try{
			ClassPool classPool = getClassPool(loader);
			CtClass ctClass = classPool.makeClass(new ByteArrayInputStream(classfileBuffer));
			//check @Service and @TransactionManagement
			Service service = null;
			TransactionManagement transactionManagement = null;
			
			for (Object o : ctClass.getAvailableAnnotations()){
				if (Service.class.isInstance(o)){
					service = (Service)o;
				}else if (TransactionManagement.class.isInstance(o)){
					transactionManagement = (TransactionManagement)o;
				}
			}
			
			if (service != null && transactionManagement != null && transactionManagement.value() == TransactionManagementType.CONTAINER ){
				loggerService.debug("find container managed service {0}", className);
				addRef(classPool, ctClass, "javax.transaction.TransactionManager", "_t_x_m_");
				for (CtMethod m : ctClass.getMethods()){
					
					int modifier = m.getModifiers();
					String mc = m.getDeclaringClass().getName();
					if (mc.startsWith("java.") || Modifier.isFinal(modifier) || Modifier.isStatic(modifier) ||  Modifier.isProtected(modifier)  || Modifier.isPrivate(modifier) || m.getMethodInfo().isConstructor() ) {
						continue;
					}
					TransactionAttribute txAttr = getAnnotation(m, TransactionAttribute.class);
					TransactionAttributeType txType = null;
					if (txAttr != null){
						txType = txAttr.value();
					}else{
						txType = TransactionAttributeType.REQUIRED;
					}
					if (txType == TransactionAttributeType.SUPPORTS){
						continue;
					}
					classPool.importPackage(WeaverService.class.getPackage().getName());

					String newName = m.getName() + "_0_";
					boolean isVoid = m.getReturnType() == CtPrimitiveType.voidType;
					CtMethod nm = CtNewMethod.copy(m, newName , ctClass, null);
					ctClass.addMethod(nm);
					StringBuilder code = new StringBuilder();
					
					switch (txType) {
						case MANDATORY :
							code.append("{TransactionLocalInterceptor _t_l_i_ = new TransactionLocalInterceptor(_t_x_m_,").append("\"MANDATORY\");\n");
							break;
						case NEVER :
							code.append("{TransactionLocalInterceptor _t_l_i_ = new TransactionLocalInterceptor(_t_x_m_,").append("\"NEVER\");\n");
							break;
						case NOT_SUPPORTED :
							code.append("{TransactionLocalInterceptor _t_l_i_ = new TransactionLocalInterceptor(_t_x_m_,").append("\"NOT_SUPPORTED\");\n");
							break;
						case REQUIRED :
							code.append("{TransactionLocalInterceptor _t_l_i_ = new TransactionLocalInterceptor(_t_x_m_,").append("\"REQUIRED\");\n");
							break;
						case REQUIRES_NEW :
							code.append("{TransactionLocalInterceptor _t_l_i_ = new TransactionLocalInterceptor(_t_x_m_,").append("\"REQUIRES_NEW\");\n");
							break;
						case SUPPORTS :
							code.append("{TransactionLocalInterceptor _t_l_i_ = new TransactionLocalInterceptor(_t_x_m_,").append("\"SUPPORTS\");\n");
							break;
						default :
							code.append("{TransactionLocalInterceptor _t_l_i_ = new TransactionLocalInterceptor(_t_x_m_);\n");;
					}
					
					
					code.append("try{ ").append(isVoid ? "" : "return ").append(newName +"($$);}").append("catch(Throwable $e) ");
					CtClass[] ets = m.getExceptionTypes();
					final StringBuffer catchCode = new StringBuffer("_t_l_i_.catchException($e);");
					for (CtClass et : ets){
						catchCode.append("if ( $e instanceof ").append(et.getName()).append(") {\n")
						.append("throw (").append(et.getName()).append(")$e;")
						.append("}\n");
					}
					catchCode.append("throw TransactionLocalInterceptor.wrapRuntimeException($e);");
					code.append("{").append(catchCode).append("}");
					code.append("finally{").append("_t_l_i_.complete();").append("}}");
//					m.insertAfter("txh$tmp.complete();", true);
					m.setBody(code.toString());
				}
				
				//remove annotations
//				
//				for (MethodInfo mi : (List<MethodInfo>)ctClass.getClassFile().getMethods()) {
//					if (mi.getName().endsWith("_0_")){
//						MethodInfo nmi = ctClass.getClassFile().getMethod(mi.getName().substring(0, mi.getName().length() - 3));
//						List<AttributeInfo> norg = (List<AttributeInfo>) nmi.getAttributes();
//						List<AttributeInfo> org = (List<AttributeInfo>)mi.getAttributes();
//						List<AttributeInfo> tmp =  new ArrayList<AttributeInfo>(org);
//						for (AttributeInfo ai : norg){
//							if (ai instanceof AnnotationsAttribute) {
//								System.out.println(ai);
//							}
//						}
//						int i = 0;
//						for (AttributeInfo ai : tmp){
//							if (ai instanceof AnnotationsAttribute) {
//								org.remove(i);
//							}
//							i++;
//						}
//						
//						
//						
//					}
//				}
				
				rt = ctClass.toBytecode();
			}
			
		}catch (Throwable e) {
			loggerService.error("transform class {0} error {1} ", className, e);
		}
		return rt;
		
	}

	@Activate
	public void activate() {
		PreMainAgent.addTransformer(this);
	}
	
	@Deactivate
	public void deactivate() {
		PreMainAgent.removeTransformer(this);
		classPoolMap.clear();
	}
	
	
}
